import os
import numpy as np
from sklearn.model_selection import train_test_split
from .feature_extractor import extract_features


def load_dataset(dataset_path, max_pad_len=100):
    X = []
    y = []

    # 遍历数据集文件夹，每个actor文件夹对应一个说话者
    for actor_folder in os.listdir(dataset_path):
        actor_path = os.path.join(dataset_path, actor_folder)

        # 遍历每个actor的音频文件
        if os.path.isdir(actor_path):
            for file_name in os.listdir(actor_path):
                if file_name.endswith('.wav'):
                    file_path = os.path.join(actor_path, file_name)
                    features = extract_features(file_path, max_pad_len=max_pad_len)
                    X.append(features)

                    # 获取情绪标签，假设文件名中包含情绪标签信息
                    label = int(file_name.split('-')[2]) - 1  # 情绪标签（0到7）
                    y.append(label)

    X = np.array(X)  # (1440, 100, 13)
    y = np.array(y)  # (1440,)

    return train_test_split(X, y, test_size=0.2, random_state=42)
